import cv2
import numpy as np
import math
import csv
import os
import sys
import subprocess

# -------------------------- 虚拟显示设置 --------------------------
def setup_virtual_display():
    """设置虚拟显示环境"""
    # 检查是否已经有显示
    if 'DISPLAY' in os.environ and os.environ['DISPLAY']:
        print(f"✅ 使用现有显示: {os.environ['DISPLAY']}")
        return None

    print("⚠️ 检测到无图形界面，尝试启动虚拟显示...")

    # 方法1: 使用pyvirtualdisplay（如果已安装）
    try:
        from pyvirtualdisplay import Display
        display = Display(visible=0, size=(1200, 800))
        display.start()
        print("✅ 使用pyvirtualdisplay启动虚拟显示成功")
        return display
    except ImportError:
        print("📝 pyvirtualdisplay未安装，尝试其他方法...")
    except Exception as e:
        print(f"❌ pyvirtualdisplay启动失败: {e}")

    # 方法2: 手动设置虚拟显示
    os.environ['DISPLAY'] = ':99'
    os.environ['QT_QPA_PLATFORM'] = 'xcb'

    # 检查Xvfb是否可用
    try:
        result = subprocess.run(['which', 'Xvfb'], capture_output=True, text=True)
        if result.returncode == 0:
            print("✅ Xvfb可用，使用DISPLAY=:99")
            return None
        else:
            print("❌ Xvfb未安装")
    except:
        print("❌ 无法检查Xvfb")

    return None

# 设置虚拟显示
virtual_display = setup_virtual_display()

# -------------------------- 全局配置与变量 --------------------------
IMAGE_PATH = "butterfly_with_ruler.jpeg"  # 替换为你的蝴蝶图像路径
EXPORT_BASE_PATH = "butterfly_measurement"
EXPORT_TXT_PATH = f"{EXPORT_BASE_PATH}.txt"
EXPORT_CSV_PATH = f"{EXPORT_BASE_PATH}.csv"

contour_points = []
is_select_complete = False
fitted_segments = []
auto_contour_points = []
is_calibrating = False
calib_points = []
PIXEL_TO_MM = 0.1  # 初始像素-毫米比例（可通过校准修改）
img_scale = 1.0    # 图像缩放比例（自适应窗口用）

# -------------------------- 核心工具函数 --------------------------
def preprocess_image(img_path, max_size=(1200, 800)):
    """图像预处理：读取、缩放、边缘检测、自动轮廓提取"""
    # 读取图像
    img = cv2.imread(img_path)
    if img is None:
        raise ValueError(f"无法读取图像：{img_path}")

    # 1. 自适应缩放
    h, w = img.shape[:2]
    scale_w = max_size[0] / w
    scale_h = max_size[1] / h
    global img_scale
    img_scale = min(scale_w, scale_h, 1.0)
    if img_scale < 1.0:
        new_w, new_h = int(w * img_scale), int(h * img_scale)
        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

    # 2. 边缘检测
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blurred, 50, 150)
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.dilate(edges, kernel, iterations=1)

    # 3. 绘制边缘到原图
    img_with_edge = img.copy()
    img_with_edge[edges != 0] = [0, 255, 0]

    # 4. 自动提取最大轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    auto_contour = []
    if contours:
        max_contour = max(contours, key=cv2.contourArea)
        epsilon = 0.002 * cv2.arcLength(max_contour, closed=False)
        auto_contour = cv2.approxPolyDP(max_contour, epsilon, closed=False)
        auto_contour = [tuple(pt[0]) for pt in auto_contour]

    return img, img_with_edge, auto_contour


def split_contour_segments(points):
    """智能分段：按曲率自动划分轮廓段（弯曲处细分，直线处粗分）"""
    if len(points) < 4:
        return [points]  # 点数量过少，不分段
    points = np.array(points)
    segments = []
    start_idx = 0
    min_seg_len = 2  # 最小段长（至少2个点）
    max_seg_len = 6  # 最大段长（最多6个点）

    # 计算每个点的曲率（用相邻3点夹角表示，夹角越小曲率越大）
    curvatures = []
    for i in range(1, len(points)-1):
        v1 = points[i] - points[i-1]
        v2 = points[i+1] - points[i]
        # 夹角余弦值（范围[-1,1]，值越小弯曲越明显）
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        curvatures.append(1 - cos_angle)  # 曲率指标（值越大弯曲越厉害）

    # 按曲率动态分段
    while start_idx < len(points) - min_seg_len:
        # 取后续3个点的平均曲率判断弯曲程度
        end_curr = min(start_idx + 3, len(curvatures))
        avg_curv = np.mean(curvatures[start_idx:end_curr]) if end_curr > start_idx else 0
        # 曲率大→细分（段长2），曲率小→粗分（段长6）
        seg_len = min_seg_len if avg_curv > 0.2 else max_seg_len
        end_idx = min(start_idx + seg_len, len(points))
        # 确保段长不小于最小值
        if end_idx - start_idx < min_seg_len:
            end_idx = start_idx + min_seg_len
        # 相邻段重叠1个点，避免拟合断裂
        segments.append(points[start_idx:end_idx].tolist())
        start_idx = end_idx - 1

    return segments


def calculate_curve_radian(sampled_points):
    """滑动窗口弧度计算：多圆拟合减少误差（返回平均弧度）"""
    if len(sampled_points) < 10:
        return 0.0  # 采样点过少，无法准确计算
    sampled_points = np.array(sampled_points)
    radian_list = []
    window_size = 8  # 滑动窗口大小（8个点拟合1个圆）

    # 滑动窗口遍历采样点
    for i in range(len(sampled_points) - window_size):
        window_pts = sampled_points[i:i+window_size].astype(np.float32)
        # 拟合外接圆
        (x0, y0), r = cv2.minEnclosingCircle(window_pts)
        if r < 2:  # 半径过小，视为直线段（无弧度）
            continue
        # 计算窗口内实际弧长
        window_arc_len = 0.0
        for j in range(1, window_size):
            dx = window_pts[j, 0] - window_pts[j-1, 0]
            dy = window_pts[j, 1] - window_pts[j-1, 1]
            window_arc_len += math.hypot(dx, dy)
        # 弧度 = 弧长 / 半径（过滤异常值）
        radian = window_arc_len / r
        if 0.1 < radian < 5:  # 排除过小（近似直线）或过大（异常拟合）的弧度
            radian_list.append(radian)

    # 返回平均弧度（无有效数据则返回0）
    return round(np.mean(radian_list) if radian_list else 0.0, 4)


def export_measurement_result():
    """多格式导出结果：TXT（汇总）+ CSV（分段详情）"""
    if not is_select_complete or not fitted_segments:
        print("❌ 未完成测量，无法导出结果！")
        return

    # 1. 计算汇总数据
    total_len_px = sum(seg["长度_像素"] for seg in fitted_segments)
    total_len_mm = total_len_px * PIXEL_TO_MM / img_scale  # 还原缩放前尺寸
    valid_radians = [seg["弧度"] for seg in fitted_segments if seg["弧度"] > 0]
    avg_radian = round(np.mean(valid_radians) if valid_radians else 0.0, 4)

    # 2. 导出TXT（汇总信息）
    with open(EXPORT_TXT_PATH, "w", encoding="utf-8") as f:
        f.write("=" * 50 + "\n")
        f.write("          蝴蝶轮廓测量结果汇总\n")
        f.write("=" * 50 + "\n")
        f.write(f"图像路径：{IMAGE_PATH}\n")
        f.write(f"像素-毫米比例：1px = {PIXEL_TO_MM:.4f}mm（缩放比例：{img_scale:.2f}）\n")
        f.write(f"总轮廓点数：{len(contour_points)}个\n")
        f.write(f"分段数量：{len(fitted_segments)}段\n")
        f.write(f"总长度（像素）：{total_len_px:.2f}px\n")
        f.write(f"总长度（毫米）：{total_len_mm:.2f}mm\n")
        f.write(f"有效弧度段数：{len(valid_radians)}段\n")
        f.write(f"平均弧度：{avg_radian:.4f}（无量纲）\n")
        f.write("=" * 50 + "\n")

    # 3. 导出CSV（分段详细信息）
    with open(EXPORT_CSV_PATH, "w", newline="", encoding="utf-8") as f:
        fieldnames = ["段号", "类型", "起点X", "起点Y", "终点X", "终点Y",
                      "长度_像素", "长度_毫米", "弧度", "采样点数量"]
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for seg in fitted_segments:
            writer.writerow(seg)

    print(f"✅ 结果已导出！")
    print(f"   TXT汇总：{EXPORT_TXT_PATH}")
    print(f"   CSV详情：{EXPORT_CSV_PATH}")


def mouse_callback(event, x, y, flags, param):
    """鼠标回调：处理选点（添加/删除/移动）、校准选点、实时预览"""
    global contour_points, is_select_complete, is_calibrating, calib_points, PIXEL_TO_MM
    img_display = param["img_display"]
    img_original = param["img_original"]
    h, w = img_display.shape[:2]

    # -------------------------- 1. 校准模式逻辑 --------------------------
    if is_calibrating:
        if event == cv2.EVENT_LBUTTONDOWN:
            calib_points.append((x, y))
            # 绘制校准点
            cv2.circle(img_display, (x, y), 4, (255, 255, 0), -1)
            cv2.putText(img_display, f"校准点{len(calib_points)}", (x+6, y-6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 1)
            # 选完2个点，计算比例
            if len(calib_points) == 2:
                # 计算校准点像素距离（还原缩放前尺寸）
                dx = calib_points[1][0] - calib_points[0][0]
                dy = calib_points[1][1] - calib_points[0][1]
                calib_px_len = math.hypot(dx, dy) / img_scale
                # 输入实际长度
                actual_len = input("请输入校准点间实际长度（单位：毫米）：")
                try:
                    actual_len = float(actual_len)
                    PIXEL_TO_MM = actual_len / calib_px_len
                    print(f"✅ 校准完成！1px = {PIXEL_TO_MM:.4f}mm")
                except ValueError:
                    print("❌ 输入错误（需为数字），校准失败！")
                # 重置校准状态
                is_calibrating = False
                calib_points = []
                # 刷新显示（重新加载边缘图）
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
        return

    # -------------------------- 2. 选点未完成时的操作 --------------------------
    if not is_select_complete:
        # （1）左键单击：添加选点
        if event == cv2.EVENT_LBUTTONDOWN:
            contour_points.append((x, y))
            # 绘制选点和序号
            cv2.circle(img_display, (x, y), 4, (0, 0, 255), -1)
            cv2.putText(img_display, str(len(contour_points)), (x+6, y-6),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
            # 绘制临时轮廓线
            if len(contour_points) >= 2:
                cv2.polylines(img_display, [np.array(contour_points)],
                              False, (255, 0, 0), 1, cv2.LINE_AA)
            # 实时显示临时长度
            temp_len_px = 0.0
            for i in range(1, len(contour_points)):
                dx = contour_points[i][0] - contour_points[i-1][0]
                dy = contour_points[i][1] - contour_points[i-1][1]
                temp_len_px += math.hypot(dx, dy)
            temp_len_mm = temp_len_px * PIXEL_TO_MM / img_scale
            # 绘制长度显示框（避免遮挡）
            cv2.rectangle(img_display, (10, h-40), (350, h-10), (0, 0, 0), -1)
            cv2.putText(img_display, f"临时长度：{temp_len_px:.1f}px ({temp_len_mm:.2f}mm)",
                        (15, h-20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1)

        # （2）右键单击：删除最近的选点（距离<10像素）
        elif event == cv2.EVENT_RBUTTONDOWN and contour_points:
            min_dist = 1000
            del_idx = -1
            # 找到最近的选点
            for i, (px, py) in enumerate(contour_points):
                dist = math.hypot(x - px, y - py)
                if dist < min_dist and dist < 10:
                    min_dist = dist
                    del_idx = i
            # 删除并刷新显示
            if del_idx != -1:
                del contour_points[del_idx]
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
                # 重新绘制剩余选点和轮廓
                for i, (px, py) in enumerate(contour_points, 1):
                    cv2.circle(img_display, (px, py), 4, (0, 0, 255), -1)
                    cv2.putText(img_display, str(i), (px+6, py-6),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                if len(contour_points) >= 2:
                    cv2.polylines(img_display, [np.array(contour_points)],
                                  False, (255, 0, 0), 1, cv2.LINE_AA)
                print(f"🗑️ 删除选点{del_idx+1}")

        # （3）按住左键：移动选点
        elif flags == cv2.EVENT_FLAG_LBUTTON and contour_points:
            min_dist = 1000
            move_idx = -1
            # 找到要移动的选点
            for i, (px, py) in enumerate(contour_points):
                dist = math.hypot(x - px, y - py)
                if dist < min_dist and dist < 10:
                    min_dist = dist
                    move_idx = i
            # 移动并刷新显示
            if move_idx != -1:
                contour_points[move_idx] = (x, y)
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
                # 重新绘制选点和轮廓
                for i, (px, py) in enumerate(contour_points, 1):
                    cv2.circle(img_display, (px, py), 4, (0, 0, 255), -1)
                    cv2.putText(img_display, str(i), (px+6, py-6),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                if len(contour_points) >= 2:
                    cv2.polylines(img_display, [np.array(contour_points)],
                                  False, (255, 0, 0), 1, cv2.LINE_AA)

    # -------------------------- 3. 双击结束选点 --------------------------
    if event == cv2.EVENT_LBUTTONDBLCLK and not is_select_complete and len(contour_points) >= 3:
        is_select_complete = True
        # 1. 智能分段
        global fitted_segments
        segments = split_contour_segments(contour_points)
        # 2. 计算每段的长度、弧度
        fitted_segments = []
        for seg_idx, seg_points in enumerate(segments, 1):
            # 计算段长度（像素，还原缩放前尺寸）
            seg_len_px = 0.0
            for i in range(1, len(seg_points)):
                dx = seg_points[i][0] - seg_points[i-1][0]
                dy = seg_points[i][1] - seg_points[i-1][1]
                seg_len_px += math.hypot(dx, dy)
            seg_len_px /= img_scale  # 还原缩放
            seg_len_mm = seg_len_px * PIXEL_TO_MM
            # 采样点（用于弧度计算，避免点过少）
            sampled_pts = []
            for i in range(len(seg_points)):
                # 线性插值增加采样点（每2像素插1个点）
                if i < len(seg_points) - 1:
                    x1, y1 = seg_points[i]
                    x2, y2 = seg_points[i+1]
                    dist = math.hypot(x2-x1, y2-y1)
                    if dist > 2:
                        steps = int(dist / 2)
                        for step in range(steps + 1):
                            tx = x1 + (x2 - x1) * step / steps
                            ty = y1 + (y2 - y1) * step / steps
                            sampled_pts.append((tx, ty))
                else:
                    sampled_pts.append(seg_points[i])
            # 计算弧度
            radian = calculate_curve_radian(sampled_pts)
            # 记录段信息
            fitted_segments.append({
                "段号": seg_idx,
                "类型": "曲线段" if radian > 0 else "直线段",
                "起点X": seg_points[0][0] / img_scale,
                "起点Y": seg_points[0][1] / img_scale,
                "终点X": seg_points[-1][0] / img_scale,
                "终点Y": seg_points[-1][1] / img_scale,
                "长度_像素": round(seg_len_px, 2),
                "长度_毫米": round(seg_len_mm, 2),
                "弧度": radian,
                "采样点数量": len(sampled_pts)
            })
        # 显示测量完成信息
        cv2.rectangle(img_display, (w//2 - 200, h//2 - 30), (w//2 + 200, h//2 + 10), (0, 255, 0), -1)
        cv2.putText(img_display, "✅ 测量完成！按 S 导出结果，按 R 重新选点",
                    (w//2 - 190, h//2 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 1)
        print(f"✅ 选点完成（{len(contour_points)}个点），已自动分段为{len(fitted_segments)}段")


def safe_set_mouse_callback(window_name, callback, param):
    """安全设置鼠标回调，处理窗口创建失败的情况"""
    try:
        # 先尝试创建窗口
        cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
        cv2.waitKey(100)

        # 检查窗口是否真的创建成功
        try:
            if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) >= 0:
                cv2.setMouseCallback(window_name, callback, param)
                print("✅ 窗口创建成功，已设置鼠标回调")
                return True
        except:
            pass

        print("❌ 窗口创建但无法访问，可能无图形支持")
        return False

    except Exception as e:
        print(f"❌ 窗口创建失败: {e}")
        return False

# -------------------------- 主程序入口 --------------------------
def main():
    global contour_points, is_select_complete, is_calibrating, auto_contour_points

    # 1. 初始化：图像预处理
    try:
        img_original, img_display, auto_contour_points = preprocess_image(IMAGE_PATH)
        print("✅ 图像预处理成功")
    except Exception as e:
        print(f"❌ 程序初始化失败：{e}")
        return

    h, w = img_display.shape[:2]

    # 2. 尝试创建窗口和设置回调
    window_created = safe_set_mouse_callback("蝴蝶轮廓测量工具", mouse_callback, {
        "img_display": img_display,
        "img_original": img_original
    })

    if not window_created:
        print("""
❌ 无法创建图形窗口，可能的原因：
1. 系统没有图形界面支持
2. OpenCV 编译时没有GUI支持
3. 缺少必要的图形库

💡 解决方案：
1. 在有图形界面的环境中运行
2. 或者使用命令行模式（需要修改代码）
3. 安装图形支持: sudo apt install libgtk2.0-0 libgtk-3-0
""")
        return

    # 3. 显示操作提示
    print("=" * 60)
    print(" 蝴蝶轮廓测量工具 - 操作说明")
    print("=" * 60)
    print("1. 基础操作：")
    print(" - 左键单击：添加选点")
    print(" - 右键单击：删除选点")
    print(" - 按住左键：移动选点")
    print(" - 双击左键：结束选点")
    print("2. 快捷功能：")
    print(f" - 按 A 键：加载自动轮廓")
    print(" - 按 C 键：进入校准模式")
    print(" - 按 S 键：导出结果")
    print(" - 按 R 键：重置选点")
    print(" - 按 ESC 键：退出程序")
    print("=" * 60)

    # 4. 主循环
    try:
        while True:
            cv2.imshow("蝴蝶轮廓测量工具", img_display)
            key = cv2.waitKey(1) & 0xFF

            if key == ord('a'):
                if not auto_contour_points:
                    print("❌ 未检测到有效轮廓，无法加载自动选点！")
                    continue
                # 重置选点并加载自动轮廓
                contour_points = auto_contour_points.copy()
                is_select_complete = False
                # 刷新显示：绘制自动选点
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
                for i, (x, y) in enumerate(contour_points, 1):
                    cv2.circle(img_display, (x, y), 4, (0, 0, 255), -1)
                    cv2.putText(img_display, str(i), (x+6, y-6),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                if len(contour_points) >= 2:
                    cv2.polylines(img_display, [np.array(contour_points)],
                                  False, (255, 0, 0), 1, cv2.LINE_AA)
                print(f"✅ 已加载自动轮廓（{len(contour_points)}个点），可手动微调")

            elif key == ord('c'):
                if is_select_complete:
                    print("❌ 已完成测量，请重置选点（按R）后再校准！")
                    continue
                is_calibrating = True
                # 显示校准提示
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
                cv2.putText(img_display, "📏 校准模式：单击左键选2个校准点（如尺子两端）",
                            (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
                print("📏 进入校准模式：请选2个校准点，之后输入实际长度（毫米）")

            elif key == ord('s'):
                export_measurement_result()

            elif key == ord('r'):
                contour_points = []
                is_select_complete = False
                is_calibrating = False
                calib_points = []
                _, img_display[:], _ = preprocess_image(IMAGE_PATH)
                print("🔄 已重置选点，可重新开始测量")

            elif key == 27:
                break

    except KeyboardInterrupt:
        print("👋 程序被用户中断")
    except Exception as e:
        print(f"❌ 程序运行错误: {e}")
    finally:
        cv2.destroyAllWindows()
        if virtual_display:
            virtual_display.stop()

if __name__ == "__main__":
    main()